import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d  import Axes3D
class PSO(object):
    def __init__(self, func, n_partciles=30, dims=2, iters=200, w=0.9, lr=(0.5, 1.5), rangep=(-2, 2), rangev=(-0.5, 0.5)):
        self.func = func
        self.n_partciles = n_partciles  # 种群大小
        self.dims = dims  # 粒子维度
        self.iters = iters  # 迭代次数
        self.w = w
        self.lr = lr  # 个体粒子的搜索经验，全体的搜索经验
        self.rangep = rangep  # 粒子搜索的限制范围
        self.rangev = rangev  # 粒子速度的限制范围
    def initSwarm(self):
        # 初始化粒子位置与速度，计算初始化下的函数适应度
        particles = np.zeros((self.n_partciles, self.dims))
        v = np.zeros((self.n_partciles, self.dims))
        fitness = np.zeros(self.n_partciles)
        for i in range(self.n_partciles):
            particles[i] = [(np.random.rand() - 0.5) * self.rangep[1] for i in range(self.dims)]
            v[i] = [(np.random.rand() - 0.5) * self.rangep[1] for i in range(self.dims)]
            fitness[i] = self.func(particles[i])
        # 初始化pbest, gbest
        gbest, gbestfitness = particles[fitness.argmin()].copy(), fitness.min()
        pbest, pbestfitness = particles.copy(), fitness.copy()
        return particles, v, fitness, gbest, gbestfitness, pbest, pbestfitness
    def run(self):
        particles, v, fitness, gbest, gbestfitness, pbest, pbestfitness = self.initSwarm()
        # 迭代循环
        result = np.zeros(self.iters)
        fig = plt.figure()
        plt.ion()
        X = np.arange(-2, 2, 0.05)
        Y = np.arange(-2, 2, 0.05)
        X, Y = np.meshgrid(X, Y)
        Z = f(np.array([X, Y]))
        cs = plt.contourf(X, Y, Z)
        fig.colorbar(cs)
        for i in range(self.iters):
            # 速度更新
            for j in range(self.n_partciles):
                v[j] = self.w * v[j] + self.lr[0] * np.random.rand() * (pbest[j] - particles[j]) \
                       + self.lr[1] * np.random.rand() * (gbest - particles[j])
            v[v < self.rangev[0]] = self.rangev[0]
            v[v > self.rangev[1]] = self.rangev[1]
            # 位置更新
            for j in range(self.n_partciles):
                particles[j] += v[j]
            particles[particles < self.rangep[0]] = self.rangep[0]
            particles[particles > self.rangep[1]] = self.rangep[1]
            s = plt.scatter(particles[:, 0], particles[:, 1], c='orange')
            plt.pause(0.4)
            s.remove()
            # 适应度更新
            for j in range(self.n_partciles):
                fitness[j] = self.func(particles[j])
            for j in range(self.n_partciles):
                if fitness[j] < pbestfitness[j]:
                    pbestfitness[j] = fitness[j]
                    pbest[j] = particles[j].copy()
            if pbestfitness.min() < gbestfitness:
                gbestfitness = pbestfitness.min()
                gbest = particles[pbestfitness.argmin()].copy()
            result[i] = gbestfitness
        plt.ioff()
        # 显示图形
        plt.scatter(gbest[0], gbest[1], c='white')
        plt.show()
        return result
def f(x):
    return np.sin(np.sqrt(np.add.reduce(np.power(x, 2))))/np.sqrt(np.add.reduce(np.power(x, 2))) + \
           np.exp(np.add.reduce(np.cos(2*np.pi*x))/2) - 2.71289
pso = PSO(f)
print(pso.run())